from net.models import *
from utils.utils import *
from utils.datasets import *
import importlib
import os
import time
import datetime
from PIL import Image
import torch

import matplotlib.pyplot as plt
import matplotlib.patches as patches

if __name__ == "__main__":
    device = select_device()
    config = importlib.import_module("param").DETECT_PARAMS
    os.makedirs(config["out_path"], exist_ok=True)

    # Set up model
    model = DarknetModel(config).to(device)

    if config["weights_path"].endswith(".weights"):
        # Load darknet weights
        model.load_darknet_weights(config["weights_path"])
    else:
        # Load checkpoint weights
        chkpt = torch.load(config["weights_path"], map_location=device)
        model.load_state_dict(chkpt['model'])
        del chkpt

    classes = load_classes(config["classname_file"])  # Extracts class labels from file
    samples = LoadImagesAndVideos(config["detect_path"], img_size=config["img_size"]) # Load samples

    imgs = []  # Stores image paths
    img_detections = []  # Stores detections for each image index

    print("\nPerforming object detection:")
    prev_time = time.time()
    model.eval()  # Set in evaluation mode
    for batch_i, (img_paths, input_imgs, image, cvcap) in enumerate(samples):
        # Configure input
        input_imgs = input_imgs.to(device)
        # Get detections
        with torch.no_grad():
            detections = model(input_imgs)
            detections = non_max_suppression(detections, config["conf_thrs"], config["nms_thrs"])

        # Log progress
        current_time = time.time()
        inference_time = datetime.timedelta(seconds=current_time - prev_time)
        prev_time = current_time
        print("\tInference Time: %s" % (inference_time))

        # Save image and detections
        # imgs.extend(img_paths)
        imgs.append(img_paths)
        img_detections.extend(detections)

    print("\nSaving images:")
    colors = []
    for i in range(config["yolo"]["num_class"]):
        color = [random.randint(0, 255) for _ in range(3)]
        colors.append(color)

    # Iterate through images and save plot of detections
    for img_i, (path, detections) in enumerate(zip(imgs, img_detections)):

        print("(%d) Image: '%s'" % (img_i, path))

        img = cv2.imread(path, cv2.IMREAD_COLOR)

        # Draw bounding boxes and labels of detections
        if detections is not None:
            # Rescale boxes to original image
            detections = rescale_boxes(detections, config["img_size"], img.shape[:2])
            unique_labels = detections[:, -1].cpu().unique()
            n_cls_preds = len(unique_labels)
            for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
                print("\t+ Label: %s, Conf: %.5f" % (classes[int(cls_pred)], cls_conf.item()))
                plot_one_box(x1, y1, x2, y2, img, color=colors[int(cls_pred)], label=None)
                # plot_one_box(x1, y1, x2, y2, img, color=colors[int(cls_pred)], label=classes[int(cls_pred)])

        # Save generated image with detections
        filename = path.split("/")[-1].split(".")[0]
        cv2.imwrite(config["out_path"] + f"/{filename}.png", img)
